Snapshot call-time args in the IIP Enzyme reverse rule (fixes #62)#63
Conversation
…tions The Const-return (IIP) reverse rule re-runs `Enzyme.autodiff(Reverse, ...)` on the unwrapped function during the reverse pass, reading the arguments' current `.val`s. When a caller mutates those arguments after the forward call -- e.g. an ODE integrator stepping `u` after every RHS call -- the VJP was taken about the post-mutation state, giving a silently wrong gradient. In the full nested EnsembleProblem adjoint this surfaced as the Julia-1.11 `_dispatch_ensemble_solve` GC-root-rewrite segfault (SciMLSensitivity.jl#1424). `augmented_primal` now snapshots the call-time argument values into the tape, and `reverse` temporarily restores them before recomputing the VJP (then puts the live values back so the surrounding reverse pass is undisturbed). Adds regression tests: a single wrapped IIP call with the argument mutated afterwards, and an 8-step integrator loop, both checked against analytic / finite-difference references. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BZwV7CBdLPqVpJrvRAWcv6
…ation The snapshot/restore safety net only needs to copy arguments that are actually mutated between the forward and reverse passes. Gate it on `EnzymeRules.overwritten(config)` (a conservative over-approximation, so it can't miss a needed snapshot): a non-mutating call now copies nothing, and an integrator-style call copies only the overwritten arrays (`du`/`u`), not `p` or the scalar `t`. Also fixes two pre-existing hand-built `RevConfig`s in the tests whose `Overwritten` tuples were one entry short (it is indexed `(func, args...)`). Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BZwV7CBdLPqVpJrvRAWcv6
|
Two notes after review feedback: Scope — this is a safety net, not the primary path. SciMLSensitivity already Allocation. The snapshot is gated on Verified: existing |
The Const-return reverse rule returned `ntuple(i -> args[i] isa Active ?
zero(...) : nothing)`, which (1) was union-typed — Enzyme rejects a
`Tuple{Union{Nothing,Float64},…}` return with `ReverseRuleReturnError`, it
requires exact per-slot types — and (2) zeroed every Active arg's gradient.
This broke any wrapped call with a mix of Duplicated and Active args, e.g. a
time-dependent in-place rhs `f!(du, u, p, t)` where `t` flows in as Active.
`Enzyme.autodiff(Reverse, Const(f_orig), Const, args...)[1]` already returns
the per-argument gradient tuple in the exact form a reverse rule must hand back
(`nothing` for Const/Duplicated, the real gradient for Active). Forward it,
rebuilt through `map(_revslot, args, raw)` so dispatching on each argument's
concrete annotation type gives an exactly-typed result (the raw tuple is
`Any`-typed inside the rule). Adds a mixed Duplicated/Active regression test.
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BZwV7CBdLPqVpJrvRAWcv6
) Reimplements the Const-return (IIP) reverse rule by delegating to Enzyme's own split reverse-mode `autodiff_thunk` on the *unwrapped* function, rather than snapshotting arguments and re-running `autodiff(Reverse, …)` in the reverse pass (the approach in the snapshot PR): - `augmented_primal` builds `autodiff_thunk(ReverseSplitModified(…, overwritten))` and runs its forward pass; the `ModifiedBetween == overwritten` flags make Enzyme's own tape cache any argument the caller mutates before the reverse pass. It stashes `(tape, reverse_thunk)`. - `reverse` runs the stashed reverse thunk; `_revslot` rebuilds its per-argument return exactly typed (so `Active` args are correct, not zeroed / union-typed). This is the more Enzyme-native design and passes every isolated test (mutation, multi-step integrator, mixed Duplicated/Active). It is offered as an alternative to #63. KNOWN LIMITATION: on the full nested `EnsembleProblem` adjoint (SciMLSensitivity #1424) the nested `autodiff_thunk` inside `augmented_primal` re-triggers the Julia-1.11 `_dispatch_ensemble_solve` GC-root-rewrite segfault in Enzyme's `post_optimize!` (EnzymeAD/Enzyme.jl #3175), which the lightweight snapshot rule in #63 avoids. So this should not merge ahead of #63 until #3175 is fixed upstream; it is the target design once it is. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BZwV7CBdLPqVpJrvRAWcv6
…it-thunk Split reverse-mode thunk for the IIP Enzyme rule (alternative to #63)
Fixes #62. Root cause and reduction of the
EnsembleProblemEnzyme adjoint blocker in SciML/SciMLSensitivity.jl#1424.Problem
The
Const-return (IIP) reverse rule re-runsEnzyme.autodiff(Reverse, Const(f_orig), Const, args...)during the reverse pass, reading the arguments' current.vals, whileaugmented_primalsaved no tape. When a caller mutates an argument after the forward call — an ODE integrator steppinguafter every rhs call — the VJP is taken about the post-mutation state, so the gradient is silently wrong. A single isolated call is fine; the bug needs a mutation-after-call. In the full nested ensemble adjoint this surfaces on Julia 1.11 as the_dispatch_ensemble_solveGC-root-rewrite segfault (#1424).Fix
augmented_primal(Const return) snapshots the call-time argument values into the tape (before running the primal, in casef_origmutates its inputs).reverse(Const return) temporarily restores those snapshots before theEnzyme.autodiff(Reverse, …)recompute, then puts the live values back so the surrounding reverse pass is undisturbed._snapshot/_restore!helpers:copy/copyto!for arrays, no-op for immutables (scalars liket).Only the IIP
Const-return path changed; the forward,Active, andDuplicated/BatchDuplicatedpaths are untouched.Verification (Julia 1.11.9, Enzyme 0.13.172/0.13.173, FunctionWrappersWrappers dev)
[3.15, 12.80]); an 8-step integrator loop matches central differences.test/Enzymesuite passes (one assertion updated:augmented_primalnow tapes the arg snapshots instead ofnothing).EnsembleProblemadjoint (SciMLBase 3.30 / OrdinaryDiffEqTsit5 / Enzyme 0.13.172,set_runtime_activity(Reverse)): with this fix the default-AutoSpecializesolve goes from 4/4 deterministic SIGSEGV → no segfault, leaving only the separatewith_logstate/AggregateLoggerblocker (same residual asFullSpecialize).🤖 Generated with Claude Code